function adjustedSurvivalCurve(covariates,survivalTime,groups,censorvec,legendLabels,adjustedOnly,stratified)


%%%%%
% This function plots adjusted (and unadjusted as well if desired) survival curves. 
% The corrected group prognosis method is used for adjusting the curves.
%
% References:
% 
% [1] Nieto and Coresh, Adjusting survival curves for confounders: a review
% and a new method, American Journal of Epidemiology, 143(10):1059-1068,
% 1996.
%
% [2] Ghali, Quan, Brant, et al., Comparison of 2 methods for calculating
% adjusted survival curves from proportional hazards models, Journal of
% American medical association, 286(12):1494-1497, 2001.
%
% Inputs:
% covariates: 2D matrix containing covariates used for adjustment; rows are individuals and columns are covariates
% survivalTime - 1D vector containing survival time, should correspond to the rows in covariates
% groups - 1D vector indicating different groups by distinct numbers, should correspond to the elements in survivalTime
% censorvec - 1D boolean vector indicating right-censored data; 0 = not censored, 1 = right-censored. Should correspond to the elements in survivalTime
% legendLabels - cell array of strings containing legend labels in ascending order of group labels
% adjustedOnly - boolean that indicates whether only adjusted survival curves are to be plotted; false = both adjusted and unadjusted, true = only adjusted
% stratified - boolean that indicates whether to build separate Cox models for different groups; false = one model for all, true = separate models
% 
% Written by Joon Lee, Mar. 11, 2011
% Revised by Joon Lee, July 26, 2011
%%%%%


% initial preparation
groups= +groups;
allGroups=sort(unique(groups));
Ng=length(allGroups);
maxT=max(survivalTime);
colors={'b' 'r' 'g' 'c' 'm' 'y' 'k' 'w'};
allLegendLabels={};

% build Cox regression models
if stratified
    b=cell(1,Ng);
    H=cell(1,Ng);
    for i=1:Ng
        idx= groups==allGroups(i);
        [b{i},logl,H{i}]=coxphfit(covariates(idx,:),survivalTime(idx),'censoring',censorvec(idx),'baseline',0);
    end
else
    [b,logl,H]=coxphfit(covariates,survivalTime,'censoring',censorvec,'baseline',0);    
end

% plot survival curves
figure
hold on

for i=1:Ng    
    
    idx= groups==allGroups(i);  
    
    if adjustedOnly        
        allLegendLabels=[allLegendLabels ['Adjusted ' legendLabels{i}]];        
    else        
        % unadjusted Kaplan-Meier curves
        [f,x]=ecdf(survivalTime(idx),'censoring',censorvec(idx),'function','survivor');
        stairs(x,f,[colors{mod(i,8)} '--'],'LineWidth',2)        
        allLegendLabels=[allLegendLabels ['Unadjusted ' legendLabels{i}] ['Adjusted ' legendLabels{i}]];        
    end
    
    % adjusted survival curves using corrected group prognosis method
    if stratified
        currb=b{i};
        Sbase=exp(-H{i}(:,2));
        t=H{i}(:,1);
    else
        currb=b;
        Sbase=exp(-H(:,2));
        t=H(:,1);
    end   
    adjustment=exp(covariates(idx,:)*currb);
    St=zeros(length(Sbase),1);
    for j=1:sum(idx)
        St=St+Sbase.^adjustment(j);
    end
    stairs(t,St./sum(idx),colors{mod(i,8)},'LineWidth',2)
    
end

xlabel('Survival Time','FontSize',14)
ylabel('Survival','FontSize',14)
legend(allLegendLabels)
xlim([0 maxT])
set(gca,'FontSize',14)

